package cn.lancetop.tools.calcc;

import stardrms.lancetop.java.cvml4j.core.cv4j.Color;
import stardrms.lancetop.java.cvml4j.core.math.Matrix;
import stardrms.lancetop.java.cvml4j.core.math.MatrixType;
import stardrms.lancetop.java.cvml4j.core.ml4j.regression.GradientDescentParameter;
import stardrms.lancetop.java.cvml4j.process.image.ImageProc;
import stardrms.lancetop.java.cvml4j.process.training.TrainProc;
import stardrms.lancetop.java.cvml4j.useful.io.OIO;
import stardrms.lancetop.java.cvml4j.useful.log.Log;
import stardrms.lancetop.java.cvml4j.useful.param.Param2;
import stardrms.lancetop.java.cvml4j.useful.thread.ConcurrentExecutor;
import stardrms.lancetop.java.cvml4j.view.Chart;
import stardrms.lancetop.java.cvml4j.view.DiscreteChart;
import stardrms.lancetop.java.cvml4j.view.FunctionChart;

import java.io.File;
import java.util.ArrayList;
import java.util.List;

public class Main {
    public static final String WORK_SPACE = "res/datas";
    public static final ImageProc imgproc = ImageProc.getInstance();
    public static final OIO oio = OIO.getInstance(new File(WORK_SPACE));

    public static final String[] tags = {
            "co",
            "so2",
            "no2",
            "o3",
            "pm2p5",
            "pm10"
    };

    public static final double[] COE = {
            0.001, 1.0, 1.0, 1.0, 1.0, 1.0
    };

    public static final int[] AQI = {
            0, 50, 100, 150, 200, 300, 400, 500
    };
    public static final int[][] LIMITS = {
            {0, 5, 10, 35, 60, 90, 120, 150},
            {0, 150, 500, 650, 800, 800, 800, 800},
            {0, 100, 200, 700, 1200, 2340, 3090, 3840},
            {0, 160, 200, 300, 400, 800, 1000, 1200},
            {0, 35, 75, 115, 150, 250, 350, 500},
            {0, 50, 150, 250, 350, 420, 500, 600}
    };

    public static double calcPoint(String tagName, double c) {
        int tagIndex = 0;
        for (tagIndex = 0; tagIndex < tags.length; tagIndex++) {
            if (tags[tagIndex].equals(tagName)) break;
        }
        int cl = 0, ch = 0, il = 0, ih = 0;
        int i;
        for (i = 1; i < LIMITS[tagIndex].length; i++) {
            if (c <= LIMITS[tagIndex][i]) {
                ch = LIMITS[tagIndex][i];
                cl = LIMITS[tagIndex][i - 1];
                ih = AQI[i];
                il = AQI[i - 1];
                break;
            }
        }
        if (i >= LIMITS[tagIndex].length) return -1;
        return (ih - il) * 1.0 / (ch - cl) * (c - cl) + il;
    }

    public static void prepareDatas() {
        int size = 1000;
        for (int i = 0; i < tags.length; i++) {
            String tagName = tags[i];
            oio.refreshData(tagName);
            double[] datas = new double[size * 2];
            int j = 0;
            int val0 = LIMITS[i][0];
            double diff = (LIMITS[i][1] - val0) * 1.0 / size;
            for (int k = 0; k < size; k++) {
                double point = calcPoint(tagName, val0 + k * diff);
                datas[j++] = point;
                datas[j++] = val0 + k * diff;
            }
            oio.writeData(Matrix.matrixFC1(size, 2, datas), tagName);
        }
    }

    public static void previewData(String dataName) {
        oio.readData(dataName, 1000, mat->{
            Log.d(mat);
            List<Param2<DiscreteChart, Matrix>> params = new ArrayList<>();
            params.add(new Param2<DiscreteChart, Matrix>(new DiscreteChart(Color.RED, Chart.VERTEX_STYLE_X, 5, false), mat));
            imgproc.plot2d(params, "Preview " + dataName);
        });
    }

    public static void normolizedExpand(String dataName) {
        oio.readData(dataName, 1000, mat->{
            Matrix x = mat.colAt(0);
            Matrix y = mat.colAt(1);
            double avg = x.clone().mean(1).at(0, 0, 0).get().doubleValue();
            double max = x.clone().max(1).at(0, 0, 0).get().doubleValue();
            double min = x.clone().min(1).at(0, 0, 0).get().doubleValue();
            Log.d(avg + ", " + max + ", " + min);
            oio.writeMatrix(Matrix.matrixFC1(1, 3, avg, max, min), "n_" + dataName);
            x.subtract(avg).dotDividedBy(max - min);
            Matrix rst = x.clone();
            rst.mergeAtRight(x.clone().dotPowerOf(2.0));
            rst.mergeAtRight(x.clone().dotPowerOf(3.0));
            rst.mergeAtRight(x.clone().dotPowerOf(4.0));
            rst.mergeAtRight(x.clone().dotPowerOf(5.0));
            rst.mergeAtRight(y);
            oio.refreshData("d_" + dataName);
            oio.writeData(rst, "d_" + dataName);
            Log.d(rst);
        });
    }

    public static void training(String dataName) {
        GradientDescentParameter param = new GradientDescentParameter();

        // set param
        param.setAlpha(0.1);
        param.setLambda(0.0);
        param.setNumIterations(1000);
        param.setHypothesisFunction((theta, X)->{
            return X.clone().times(theta);
        });
        param.setCostFunction((h, theta, X, y, lambda)->{
            int m = X.numOfRows();
            Matrix part2 = theta.subMatrix(1, theta.numOfRows() - 1, 0, theta.numOfColumns() - 1).dotPowerOf(2.0).sum(1).dotTimes(lambda);
            return (h.hTheta(theta, X).subtract(y).dotPowerOf(2.0).sum(1)).add(part2).dotDividedBy(2.0 * m);
        });
        param.setGradientFunction((h, theta, X, y, lambda)->{
            int m = X.numOfRows();
            Matrix part2 = theta.subMatrix(1, theta.numOfRows() - 1, 0, theta.numOfColumns() - 1).dotTimes(lambda);
            part2.mergeAtTop(Matrix.zeros(1, 1, MatrixType.FC1));
            return X.clone().transpose().times(h.hTheta(theta, X).subtract(y)).add(part2).dotDividedBy(m * 1.0);
        });

        TrainProc.getInstance().batchGradientDescent(new File(WORK_SPACE), "d_" + dataName, param, new ConcurrentExecutor(),
                (theta, Jhistory)->{
                    Log.d(theta);
                    oio.writeMatrix(theta, "t_" + dataName);
                    oio.refreshData("p_" + dataName);
                    oio.writeData(theta, "p_" + dataName);
                    Log.d(Jhistory);
                    Log.d("J-history" + Jhistory);
                    Log.d(String.format("J-final: %8.6f", Jhistory.at(Jhistory.numOfRows() - 1, 0, 0).get().doubleValue()));
                    for (int i = 0; i < Jhistory.numOfRows(); i++) {
                        if (Double.isNaN(Jhistory.at(i, 0, 0).get().doubleValue())) Jhistory.setAt(i, 0, Jhistory.getZaxisFactory().fill(-1));
                    }
                    List<Param2<DiscreteChart, Matrix>> params = new ArrayList<>();
                    params.add(new Param2<DiscreteChart, Matrix>(new DiscreteChart(Color.RED, Chart.VERTEX_STYLE_DOT, 3, true), Jhistory));
                    imgproc.plot2d(params, "Training " + dataName + String.format("| J-final: %8.6f", Jhistory.at(Jhistory.numOfRows() - 1, 0, 0).get().doubleValue()));
                }
        );
    }

    public static int getIndex(String dataName) {
        int tagIndex = 0;
        for (tagIndex = 0; tagIndex < tags.length; tagIndex++) {
            if (tags[tagIndex].equals(dataName)) break;
        }
        return tagIndex;
    }

    public static void testing(String dataName) {
        int tagIndex = getIndex(dataName);
        Matrix theta = oio.readMatrix("t_" + dataName); // Matrix.matrixFC1(2, 1, 0.004320, 0.695540);
        Matrix n = oio.readMatrix("n_" + dataName);
        double avg = n.at(0, 0, 0).get().doubleValue();
        double max = n.at(0, 1, 0).get().doubleValue();
        double min = n.at(0, 2, 0).get().doubleValue();
        oio.readData(dataName, 1000, data->{
            Matrix x = data.colAt(0);
            Matrix y = data.colAt(1);
            Matrix X = Matrix.ones(x.numOfRows(), 1, x.getType());
            x.subtract(avg).dotDividedBy(max - min);
            X.mergeAtRight(x);
            X.mergeAtRight(x.clone().dotPowerOf(2.0));
            X.mergeAtRight(x.clone().dotPowerOf(3.0));
            X.mergeAtRight(x.clone().dotPowerOf(4.0));
            X.mergeAtRight(x.clone().dotPowerOf(5.0));
            Matrix ans = X.times(theta);
            double cnt = ans.subtract(y).abs().lessThan(1.0 * COE[tagIndex]).sum(1).at(0, 0, 0).get().doubleValue();
            Log.d(String.format("%s: %.2f%%", dataName, cnt / x.numOfRows() * 100));
        });
    }

    public static void printEquation(String dataName) {
        Matrix theta = oio.readMatrix("t_" + dataName); // Matrix.matrixFC1(2, 1, 0.004320, 0.695540);
        Matrix n = oio.readMatrix("n_" + dataName);
        double avg = n.at(0, 0, 0).get().doubleValue();
        double max = n.at(0, 1, 0).get().doubleValue();
        double min = n.at(0, 2, 0).get().doubleValue();
        double t0 = theta.at(0, 0, 0).get().doubleValue();
        double t1 = theta.at(1, 0, 0).get().doubleValue();
        double t2 = theta.at(2, 0, 0).get().doubleValue();
        double t3 = theta.at(3, 0, 0).get().doubleValue();
        double t4 = theta.at(4, 0, 0).get().doubleValue();
        double t5 = theta.at(5, 0, 0).get().doubleValue();
        Log.d(String.format("%s: x=(x-%.2f)/(%.2f-%.2f); result=%.2f+%.2f*x+%.2f*x*x+%.2f*x*x*x+%.2f*x*x*x*x + %.2f*x*x*x*x*x;", dataName, avg, max, min, t0, t1, t2, t3, t4, t5));
        imgproc.plotFunction("equation:" + dataName, 600, 600, 15, 10.0, new FunctionChart(Color.RED, x->{
            x = (x - avg) / (max - min);
            return t0 + t1 * x + t2 * x * x + t3 * x * x * x + t4 * x * x * x * x + t5 * x * x * x * x * x;
        }));
    }

    public static void main(String[] args) {
        Log.t("Start");

//        prepareDatas();

        for (String tag : tags) {
//            previewData(tag);
//            normolizedExpand(tag);
//		    training(tag);
//            testing(tag);
//            printEquation(tag);
        }

        Log.t("End");
    }

}
